import numpy as np  
import torch  
import torch.nn as nn  
from PIL import Image  
import gradio as gr  
  
# 导入定义好的多层感知机模型结构  
from modelstruct import MLP  

# 加载预训练的模型参数  
params_path = 'mnist_mlp.pth'  
loaded_params = torch.load(params_path)  
  
# 实例化模型并加载参数  
model = MLP()  
model.load_state_dict(loaded_params)  
model.eval()  # 设置模型为评估模式  
  
# 预处理函数  
def preprocess_image(image):  
    image = Image.fromarray(image).convert('L').resize((28, 28))  
    image_array = np.array(image) / 255.0  
    return torch.from_numpy(image_array).float().unsqueeze(0)  
  
# 预测函数  
def predict_digit(image):  
    preprocessed_image = preprocess_image(image)  
    with torch.no_grad():  
        output = model(preprocessed_image)  
    predicted_digit = torch.argmax(output).item()  
    return str(predicted_digit)  
  
# 创建Gradio接口  
iface = gr.Interface(fn=predict_digit, inputs='sketchpad', outputs='text')  
  
# 启动Gradio接口  
iface.launch()